import argparse
import time
from tqdm import tqdm
from pathlib import Path
import yaml
import os

from dataset.amass import AMASS
from dataset.extract import Anim
from skeleton.skeleton import Skeleton
from model import Model

MODELS_DIR = "models"

def parse_args():
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--mode",
        type=str,
        default='train',
        help="'train'/'test'",
    )

    parser.add_argument(
        "--data-path",
        type=str,
        default=f"./data",
        help="Path to img to test",
    )
    
    parser.add_argument(
        "--load-path",
        type=str,
        default=f"",
        help="Path to checkpoint to load",
    )
    
    parser.add_argument(
        "--epoch",
        type=str,
        default="",
        help="The epoch of checkpoint to load",
    )
    
    args = parser.parse_args()

    return args

if __name__ == '__main__':
    
    # argsparse
    args = parse_args()
    
    # config
    mode = args.mode
    if mode == "train":
        conf = yaml.safe_load(open('./config/config_train.yaml','r').read())
    elif mode == "test":
        conf = yaml.safe_load(open('./config/config_test.yaml','r').read())
    
    # save by time
    stamp = time.strftime("%Y-%m-%d-%H_%M_%S", time.localtime(time.time()))
    save_path = conf["save"]["save_path"]
    model_path = Path(save_path) / stamp / 'checkpoint'
    log_path = Path(save_path) / stamp / 'log'
    conf["save"]["save_path"] =  model_path
    conf["save"]["log_path"] = log_path
    Path(model_path).mkdir(exist_ok=True, parents=True)
    Path(log_path).mkdir(exist_ok=True, parents=True)
    
    print("Loading Settings Finished.")
    
    # data
    data_path = args.data_path
    train_data = None
    valid_data = None
    test_data = None
    if mode == "train":
        train_data = AMASS(data_path, "train", window=conf["data"]["sequence_length"],
                           offset=80, flip=conf["data"]["flip"])
        valid_data = AMASS(data_path, "valid", window=conf["data"]["sequence_length"],
                           offset=80, flip=conf["data"]["flip"])
    elif mode =="test":
        test_data = AMASS(data_path, "test", window=conf["data"]["sequence_length"],
                           offset=80, flip=conf["data"]["flip"])
    print("Loading HMR Data Finished.")
      
    # Skeleton
    skeleton = Skeleton(Anim.OFFSETS, Anim.PARENTS).cuda()
    
    print("Initializing Skeleton Finished.")
    
    # model
    if mode == "train":
        net = Model(conf, skeleton, train_data, valid_data)
    elif mode == "test":
        net = Model(conf, skeleton, train_data, test_data)
    
    print("Initializing Model Finished.")
    
    if mode == "train":
        net.load_path = args.load_path
        if net.load_path!=f"":
            print(net.load_path)
            net.load(epoch=args.epoch)
        net.train()
    elif mode == "test":
        net.load_path = args.load_path
        print(net.load_path)
        net.load(epoch=args.epoch)
        net.test()